import os
os.environ["QT_X11_NO_MITSHM"] = "1"  # 必须在导入PyQt5之前设置
import sys
if hasattr(sys, 'frozen'):
    os.environ['QT_QPA_PLATFORM_PLUGIN_PATH'] = os.path.join(sys._MEIPASS, 'PyQt5', 'Qt', 'plugins')

from PyQt5.QtWidgets import (QApplication, QWidget, QVBoxLayout, QHBoxLayout, QPushButton, 
                             QLabel, QProgressBar, QFileDialog, QMessageBox, QGridLayout)
from PyQt5.QtGui import QImage, QPixmap, QFont
from PyQt5.QtCore import QThread, pyqtSignal, Qt, QMutex 
import time
import yaml
from pathlib import Path
import cv2
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))

os.environ.pop("QT_QPA_PLATFORM_PLUGIN_PATH", None)

# 定义日志路径
PATH_LOGS = Path("./results")
PATH_LOGS.mkdir(parents=True, exist_ok=True)


class GUI(QWidget):
    result_signal = pyqtSignal(dict)
    
    def __init__(self, send_pause_command, date=None, i=1):
        super().__init__()
        self.send_pause_command = send_pause_command
        
        # 面列表和统计项类型
        self.names = ["up", "down", "front", "back", "left", "right"]
        self.stats_types = ["screw", "wire", "spot"]
        
        self.labels = {}  # 存储计数标签
        self.frames = {}  # 存储每个统计项的小框
        self.result_signal.connect(self._update_result_safe)
        self.thread_mutex = QMutex()

        # 加载标准数据
        self.standard_data = self.load_standard_data()

        if date is None:
            date = time.strftime("%Y%m%d"+"_%H%M%S")
        self.date = date
        self.i = i

        # 视频和计数相关变量
        self.current_frame = None
        self.playing = False
        self.frame_delay = 40
        self.current_label = None  # 当前检测的面
        self.counting_enabled = False  # 计数开关

        # 文件存储相关变量
        self.date_folder = Path("")
        self.workpiece_folder = Path("")
        self.frame_count = 0
        
        # 初始化统计信息
        self.max_screw_info = {
            name: {
                "screw": {"count": 0, "video_time": 0.0},
                "wire": {"count": 0, "video_time": 0.0},
                "spot": {"count": 0, "video_time": 0.0}
            } for name in self.names
        }

        self.initUI()
        self.init_save_folder(self.date, self.i)

    def load_standard_data(self):
        """加载标准数据文件stander.yaml"""
        try:
            with open("src/gui/detection_gui/scripts/stander.yaml", 'r', encoding='utf-8') as f:
                return yaml.safe_load(f)
        except Exception as e:
            QMessageBox.warning(self, "加载标准数据失败", f"无法加载stander.yaml: {str(e)}")
            return {}

    def initUI(self):
        self.setWindowTitle("Camera Transform GUI")
        self.setGeometry(100, 100, 1500, 800)
        main_layout = QHBoxLayout()
        
        # 视频播放区域
        video_layout = QVBoxLayout()
        self.canvas = QLabel(self)
        self.canvas.setFixedSize(1200, 700)
        self.canvas.setStyleSheet("background-color: black;")
        video_layout.addWidget(self.canvas)
        
        # 计算区域
        calculate_layout = QVBoxLayout()
        top_part = QGridLayout()
        main_font_size = 16
        sub_font_size = 14
        box_style = "border: 2px solid blue; border-radius: 10px; padding: 10px;"
        
        # 初始化6个大框（每行2个）
        for i, name in enumerate(self.names):
            row = i // 2
            col = i % 2
            
            # 创建大框容器
            main_frame = QWidget(self)
            main_frame.setStyleSheet(box_style)
            main_frame.setMinimumHeight(200)
            main_layout_v = QVBoxLayout(main_frame)
            
            # 大框标题（面名）
            main_label = QLabel(name, self)
            main_label.setAlignment(Qt.AlignCenter)
            main_label.setStyleSheet(f"font-size: {main_font_size}px; font-weight: bold; margin: 5px;")
            main_layout_v.addWidget(main_label)
            
            # 小框容器（横向排列三个小框）
            sub_layout = QHBoxLayout()
            sub_layout.setSpacing(10)
            
            # 为每个大框添加三个小框（screw, wire, spot）
            for stat_type in self.stats_types:
                sub_frame = QWidget(self)
                sub_frame.setStyleSheet("border: 1px solid #ccc; border-radius: 5px; padding: 5px;")
                sub_vlayout = QVBoxLayout(sub_frame)
                
                # 小框标题（类型名）
                sub_label = QLabel(stat_type, self)
                sub_label.setAlignment(Qt.AlignCenter)
                sub_label.setStyleSheet(f"font-size: {sub_font_size}px; font-weight: bold; color: #555;")
                sub_vlayout.addWidget(sub_label)
                
                # 小框计数显示
                count_label = QLabel("数量: 0", self)
                count_label.setAlignment(Qt.AlignCenter)
                count_label.setStyleSheet(f"font-size: {sub_font_size}px;")
                sub_vlayout.addWidget(count_label)
                
                sub_layout.addWidget(sub_frame)
                # 存储计数标签和框架引用
                self.labels[f"{name}_{stat_type}"] = count_label
                self.frames[f"{name}_{stat_type}"] = sub_frame  # 存储小框引用
            
            main_layout_v.addLayout(sub_layout)
            top_part.addWidget(main_frame, row, col)
        
        # 下部分控制区域
        bottom_part = QVBoxLayout()

        # 合格/不合格按钮布局
        qualify_layout = QHBoxLayout()
        self.qualified_btn = QPushButton("合格", self)
        self.qualified_btn.setStyleSheet("font-size: 14px; padding: 8px;")
        self.qualified_btn.clicked.connect(self.on_qualified_clicked)
        qualify_layout.addWidget(self.qualified_btn)
        
        self.unqualified_btn = QPushButton("不合格", self)
        self.unqualified_btn.setStyleSheet("font-size: 14px; padding: 8px;")
        self.unqualified_btn.clicked.connect(self.on_unqualified_clicked)
        qualify_layout.addWidget(self.unqualified_btn)
        bottom_part.addLayout(qualify_layout)
        
        self.current_status_label = QLabel(f"当前状态: 未选择面（未计数）")
        self.current_status_label.setAlignment(Qt.AlignCenter)
        self.current_status_label.setStyleSheet("font-size: 14px; color: #333; margin: 10px;")
        bottom_part.addWidget(self.current_status_label)
        
        calculate_layout.addLayout(top_part)
        calculate_layout.addLayout(bottom_part)
        main_layout.addLayout(video_layout)
        main_layout.addLayout(calculate_layout)
        self.setLayout(main_layout)

    def init_save_folder(self, date_str, workpiece_idx):
        """初始化保存文件夹"""
        self.date_folder = Path(date_str)
        self.workpiece_folder = PATH_LOGS / self.date_folder / f"workpiece{workpiece_idx}"
        self.workpiece_folder.mkdir(parents=True, exist_ok=True)
        return self.workpiece_folder

    def on_qualified_clicked(self):
        self.send_qualify_status(True)
        self.close()

    def on_unqualified_clicked(self):
        self.send_qualify_status(False)
        self.close()

    def send_qualify_status(self, status):
        self.current_qualify_status = status
        if hasattr(self, 'ros2_node'):
            self.ros2_node.publish_qualify_status(status)

    def check_and_update_colors(self):
        """在move状态时检查所有面的计数是否符合标准（修正版）"""
        if not self.standard_data:
            return
            
        for name in self.names:
            # 确保标准数据中存在该面
            if name not in self.standard_data:
                continue
                
            for stat_type in self.stats_types:
                # 安全获取标准计数（避免键不存在的错误）
                standard_count = self.standard_data[name].get(stat_type, {}).get("count", 0)
                # 获取当前识别的计数
                current_count = self.max_screw_info[name][stat_type]["count"]
                
                # 严格按照规则判断
                if stat_type in ["screw", "wire"]:
                    # 允许±1误差
                    is_pass = abs(current_count - standard_count) <= 1
                else:  # spot
                    # 必须完全一致
                    is_pass = current_count == standard_count
                
                # 强制更新判断结果（覆盖旧值）
                frame_key = f"{name}_{stat_type}"
                if frame_key in self.frames:
                    self.frames[frame_key].setProperty("is_pass", is_pass)
                    # 调试信息：打印判断过程（方便排查）
                    print(f"面: {name}, 类型: {stat_type}, 标准: {standard_count}, 当前: {current_count}, 通过: {is_pass}")


    def switch_current_surface(self, new_input):
        self.thread_mutex.lock()
        try:
            if new_input == "move":
                # 进入move状态时，强制重新计算所有面的判断结果
                self.check_and_update_colors()
                if self.current_label is not None:
                    self.apply_color_for_surface(self.current_label)  # 刷新当前面的颜色
                self.counting_enabled = False
                self.current_status_label.setText(f"当前状态: move（计数暂停）")
            elif new_input in self.names:
                # 进入新面时，强制应用最新的颜色判断结果
                self.current_label = new_input
                self.counting_enabled = True
                self.current_status_label.setText(f"当前状态: 计数中（当前面：{new_input}）")
        finally:
            self.thread_mutex.unlock()


    def apply_color_for_surface(self, surface_name):
        """强制刷新UI颜色，确保使用最新的判断结果"""
        for stat_type in self.stats_types:
            frame_key = f"{surface_name}_{stat_type}"
            if frame_key not in self.frames:
                continue  # 跳过不存在的UI元素
            
            frame = self.frames[frame_key]
            
            # 1. 强制重新计算一次判断结果（不依赖缓存的property）
            # 获取最新的标准值和当前值
            standard_count = self.standard_data.get(surface_name, {}).get(stat_type, {}).get("count", 0)
            current_count = self.max_screw_info[surface_name][stat_type]["count"]
            
            # 重新判断（确保与check_and_update_colors逻辑一致）
            if stat_type in ["screw", "wire"]:
                is_pass = abs(current_count - standard_count) <= 1
            else:
                is_pass = current_count == standard_count
            
            # 2. 强制更新样式表（覆盖任何旧样式）
            if is_pass:
                # 绿色背景（通过）
                frame.setStyleSheet("""
                    border: 1px solid #ccc; 
                    border-radius: 5px; 
                    padding: 5px; 
                    background-color: #90EE90;  /* 浅绿色 */
                """)
            else:
                # 红色背景（不通过）
                frame.setStyleSheet("""
                    border: 1px solid #ccc; 
                    border-radius: 5px; 
                    padding: 5px; 
                    background-color: #FFB6C1;  /* 浅红色 */
                """)
            
            # 3. 强制刷新UI（解决Qt样式不更新的问题）
            frame.style().unpolish(frame)
            frame.style().polish(frame)
            frame.update()
            
            # 调试：打印最终应用的颜色
            color = "绿色（通过）" if is_pass else "红色（不通过）"
            print(f"应用 {surface_name}_{stat_type} 颜色：{color}（标准：{standard_count}，当前：{current_count}）")

        def toggle_play(self):
            if self.playing:
                self.playing = False
                if self.play_thread:
                    self.play_thread.stop()
            else:
                self.playing = True
                self.play_thread = PlayVideoThread(self)
                self.play_thread.frame_signal.connect(self.display_frame)
                self.play_thread.debug_signal.connect(self.calcu_crew)
                self.play_thread.start()

    def display_frame(self, frame):
        if frame is not None:
            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            height, width, _ = rgb_frame.shape
            canvas_width, canvas_height = 1200, 700
            ratio = min(canvas_width / width, canvas_height / height)
            new_width = int(width * ratio)
            new_height = int(height * ratio)
            resized_frame = cv2.resize(rgb_frame, (new_width, new_height))
            h, w, ch = resized_frame.shape
            bytes_per_line = ch * w
            q_img = QImage(resized_frame.data, w, h, bytes_per_line, QImage.Format_RGB888)
            pixmap = QPixmap.fromImage(q_img)
            self.canvas.setPixmap(pixmap)

    def init_video_writer(self):
        if self.workpiece_folder == Path(""):
            self.init_save_folder(time.strftime("%Y%m%d"), 1)
        video_path = self.workpiece_folder / "video.avi"
        if self.current_frame is not None:
            height, width, _ = self.current_frame.shape
            fourcc = cv2.VideoWriter_fourcc(*'XVID')
            self.out = cv2.VideoWriter(str(video_path), fourcc, 15.0, (width, height))

    def save_video(self):
        if self.playing:
            self.toggle_play()
        if hasattr(self, 'out') and self.out is not None:
            self.out.release()
        self.save_files()
        QMessageBox.information(self, "保存成功", f"视频和统计信息已保存到：\n{str(self.workpiece_folder)}")

    def _delete_existing_image(self, label, stat_type):
        if not self.workpiece_folder.exists():
            return
        target_imgs = list(self.workpiece_folder.glob(f"{label}_{stat_type}_*.png"))
        for img in target_imgs:
            try:
                img.unlink()
            except Exception as e:
                print(f"删除图片失败: {e}")

    def _save_label_image(self, label, stat_type):
        if self.current_frame is not None and self.workpiece_folder.exists():
            self._delete_existing_image(label, stat_type)
            img_name = f"{label}_{stat_type}_{self.frame_count}.png"
            img_path = self.workpiece_folder / img_name
            cv2.imwrite(str(img_path), self.current_frame)

    def save_files(self):
        if self.workpiece_folder == Path(""):
            self.init_save_folder(time.strftime("%Y%m%d"), 1)
        stats_info = {}
        for name in self.names:
            stats_info[name] = {
                "screw": self.max_screw_info[name]["screw"],
                "wire": self.max_screw_info[name]["wire"],
                "spot": self.max_screw_info[name]["spot"]
            }
        yaml_path = self.workpiece_folder / "stats_info.yaml"
        with open(yaml_path, 'w', encoding='utf-8') as f:
            yaml.dump(stats_info, f, allow_unicode=True, default_flow_style=False, sort_keys=False)

    def calcu_crew(self, data):
        """处理每个面的screw、wire、spot计数"""
        if not self.counting_enabled or not self.current_label:
            return
        
        for item in data:
            label, count = item
            # 处理当前面的screw计数
            if label == self.current_label:
                stat_type = "screw"
                stat_key = f"{self.current_label}_{stat_type}"
                if count > self.max_screw_info[self.current_label][stat_type]["count"]:
                    self.max_screw_info[self.current_label][stat_type]["count"] = count
                    video_time = self.frame_count / 15.0
                    self.max_screw_info[self.current_label][stat_type]["video_time"] = round(video_time, 2)
                    self.labels[stat_key].setText(f"数量：{count}")
                    self._save_label_image(self.current_label, stat_type)
            
            # 处理当前面的spot计数
            elif label == f"{self.current_label}_spot":
                stat_type = "spot"
                stat_key = f"{self.current_label}_{stat_type}"
                if count > self.max_screw_info[self.current_label][stat_type]["count"]:
                    self.max_screw_info[self.current_label][stat_type]["count"] = count
                    video_time = self.frame_count / 15.0
                    self.max_screw_info[self.current_label][stat_type]["video_time"] = round(video_time, 2)
                    self.labels[stat_key].setText(f"数量：{count}")
                    self._save_label_image(self.current_label, stat_type)
            
            # 处理当前面的wire计数
            elif label == f"{self.current_label}_wire":
                stat_type = "wire"
                stat_key = f"{self.current_label}_{stat_type}"
                if count > self.max_screw_info[self.current_label][stat_type]["count"]:
                    self.max_screw_info[self.current_label][stat_type]["count"] = count
                    video_time = self.frame_count / 15.0
                    self.max_screw_info[self.current_label][stat_type]["video_time"] = round(video_time, 2)
                    self.labels[stat_key].setText(f"数量：{count}")
                    self._save_label_image(self.current_label, stat_type)

    def closeEvent(self, event):
        if hasattr(self, 'play_thread') and self.play_thread:
            self.play_thread.stop()
        if hasattr(self, 'cap') and self.cap:
            self.cap.release()
        if hasattr(self, 'out') and self.out:
            self.out.release()
        self.save_video()
        event.accept()

    def set_current_result(self, result):
        """更新处理结果"""
        self.result_signal.emit(result)
    
    def _update_result_safe(self, result):
        self.current_frame = result['image']
        
        if self.counting_enabled and self.current_label:
            self.frame_count += 1
            if not hasattr(self, 'out') or self.out is None:
                self.init_video_writer()
            if hasattr(self, 'out') and self.out is not None:
                self.out.write(self.current_frame)
        
        # 提取当前帧的各类计数
        current_counts = [
            [self.current_label, result['object_count']],  # 当前面的screw计数
            [f"{self.current_label}_spot", result['spot_num']],  # 当前面的spot计数
            [f"{self.current_label}_wire", result['wire_num']]   # 当前面的wire计数
        ]
        self.calcu_crew(current_counts)
        
        if not self.playing:
            self.display_frame(cv2.resize(result['image'], (1280, 720)))


class PlayVideoThread(QThread):
    frame_signal = pyqtSignal(object)
    debug_signal = pyqtSignal(list)
    
    def __init__(self, gui):
        super().__init__()
        self.gui = gui
        self.thread_running = True
        # 测试数据
        self.debug_data = [
            [["front", 1], ["front_spot", 1], ["front_wire", 1]],
            [["front", 2], ["front_spot", 2], ["front_wire", 2]],
            [["front", 3], ["front_spot", 3], ["front_wire", 3]],
            [["front", 4], ["front_spot", 4], ["front_wire", 4]],
            [["front", 5], ["front_spot", 5], ["front_wire", 5]]
        ]
        self.debug_index = 0
    
    def run(self):
        while self.thread_running:
            if self.gui.current_frame is not None:
                if self.gui.counting_enabled:
                    self.gui.frame_count += 1
                self.debug_index = (self.debug_index + 1) % len(self.debug_data)
                self.frame_signal.emit(self.gui.current_frame)
                self.debug_signal.emit(self.debug_data[self.debug_index])
            
            time.sleep(self.gui.frame_delay / 1000)
    
    def stop(self):
        self.thread_running = False
        self.wait()


if __name__ == '__main__':
    def send_pause_command():
        pass
    app = QApplication(sys.argv)
    gui = GUI(send_pause_command)
    gui.show()
    sys.exit(app.exec_())